import torch.utils.data as data
from PIL import Image
from torchvision import transforms 
import torch
import os
import numpy as np
from torchvision import transforms as T
import utils.utils as utils
import yaml
path='./default.yaml'
with open(path, 'r') as f:
    cfg = yaml.load(f, Loader=yaml.Loader)


def readindex2prefix(filename,cams):  #eth3d

    with open(filename) as f:
        lines = [line.rstrip() for line in f.readlines()]
    return lines[cams].split()
def read_cam_file(filename):  #eth3d

    with open(filename) as f:
        lines = [line.rstrip() for line in f.readlines()]

    extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ')
    extrinsics = extrinsics.reshape((4, 4))


    intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ')
    intrinsics = intrinsics.reshape((3, 3))

    intrinsics=np.vstack([intrinsics,[0 ,0, 0]])#Vertical splicing
    intrinsics=np.hstack([intrinsics,[[0] ,[0], [0], [1]]])     #Horizontal splicing
    depth_min = float(lines[11].split()[0])
    depth_interval=float(lines[11].split()[1])
    DEPTH_NUM=float(lines[11].split()[2])
    depth_max = depth_min + 128 * depth_interval
    return intrinsics, extrinsics,depth_min,depth_max


#  delivery_area  electro    forest  playground  terrains

class Eth3D(data.Dataset):
    def __init__(self,flag='train',category='delivery_area'):
         super(Eth3D, self).__init__()
         self.height = 480
         self.width = 896
         sumnumber=int(readindex2prefix('/data/kaixindata/mvs_training/eth3d/'+category+'/cams/index2prefix.txt',0)[0])-2
        
         self.traincames=range(2,round(sumnumber*0.8))
         self.testcames=range(round(sumnumber*0.8),sumnumber)
         self.framnumber=1
         self.category=category
         self.flag=flag